//! Kafka管理器
//! 
//! 提供Kafka生产者和消费者的集中管理，与沙箱管理器集成

use std::sync::Arc;
use anyhow::{Result, anyhow};
use tracing::{info, error, warn, debug};
use tokio::sync::Mutex;
use async_trait::async_trait;

use crate::kafka::{
    KafkaConfig, KafkaProducer, KafkaProducerBuilder,
    KafkaConsumer, KafkaConsumerBuilder, MessageHandler,
    KafkaMessage, SandboxEventMessage, ExecutionTaskMessage,
    SystemMetricsMessage, ErrorReportMessage, HealthCheckMessage,
    SandboxEventType, ErrorLevel, HealthStatus,
};
use crate::types::{SandboxId, Language, SandboxState};
use crate::sandbox_manager::SandboxManager;

/// Kafka管理器
pub struct KafkaManager {
    /// 配置
    config: KafkaConfig,
    /// 生产者
    producer: Option<Arc<KafkaProducer>>,
    /// 消费者
    consumer: Option<Arc<Mutex<KafkaConsumer>>>,
    /// 沙箱管理器引用（用于处理执行任务）
    sandbox_manager: Option<Arc<SandboxManager>>,
}

impl KafkaManager {
    /// 创建新的Kafka管理器
    pub fn new(config: KafkaConfig) -> Self {
        Self {
            config,
            producer: None,
            consumer: None,
            sandbox_manager: None,
        }
    }
    
    /// 设置沙箱管理器引用
    pub fn with_sandbox_manager(mut self, sandbox_manager: Arc<SandboxManager>) -> Self {
        self.sandbox_manager = Some(sandbox_manager);
        self
    }
    
    /// 初始化生产者
    pub async fn init_producer(&mut self) -> Result<()> {
        info!("初始化Kafka生产者");
        
        let producer = KafkaProducerBuilder::new(self.config.clone())
            .client_id("koda-sandbox-manager-producer")
            .enable_idempotence(true)
            .max_in_flight_requests_per_connection(5)
            .build()?;
        
        self.producer = Some(Arc::new(producer));
        
        info!("Kafka生产者初始化完成");
        Ok(())
    }
    
    /// 初始化消费者
    pub async fn init_consumer(&mut self) -> Result<()> {
        info!("初始化Kafka消费者");
        
        let handler = SandboxMessageHandler::new(
            self.sandbox_manager.clone()
                .ok_or_else(|| anyhow!("必须先设置沙箱管理器"))?
        );
        
        let consumer = KafkaConsumerBuilder::new(self.config.clone())
            .client_id("koda-sandbox-manager-consumer")
            .handler(handler)
            .max_poll_interval_ms(300000) // 5分钟
            .build()?;
        
        self.consumer = Some(Arc::new(Mutex::new(consumer)));
        
        info!("Kafka消费者初始化完成");
        Ok(())
    }
    
    /// 启动消费者
    pub async fn start_consumer(&self) -> Result<()> {
        if let Some(consumer) = &self.consumer {
            let consumer_guard = consumer.lock().await;
            consumer_guard.consume_all_topics().await?;
        } else {
            return Err(anyhow!("消费者未初始化"));
        }
        Ok(())
    }
    
    /// 停止消费者
    pub async fn stop_consumer(&self) -> Result<()> {
        if let Some(consumer) = &self.consumer {
            let consumer_guard = consumer.lock().await;
            consumer_guard.stop().await?;
        }
        Ok(())
    }
    
    /// 发送沙箱事件
    pub async fn send_sandbox_event(
        &self,
        sandbox_id: SandboxId,
        event_type: SandboxEventType,
        sandbox_state: SandboxState,
        language: Language,
        metadata: std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<()> {
        let producer = self.producer.as_ref()
            .ok_or_else(|| anyhow!("生产者未初始化"))?;
        
        let message = SandboxEventMessage {
            header: crate::kafka::messages::MessageHeader::new(
                "sandbox_event",
                "koda-sandbox-manager"
            ),
            sandbox_id,
            event_type,
            sandbox_state,
            language,
            metadata,
        };
        
        producer.send_sandbox_event(&message).await
    }
    
    /// 发送执行任务
    pub async fn send_execution_task(
        &self,
        task_id: String,
        sandbox_id: Option<SandboxId>,
        source_code: String,
        language: Language,
        entry_point: String,
        env_vars: std::collections::HashMap<String, String>,
        stdin: String,
        priority: crate::types::Priority,
        timeout_ms: u64,
    ) -> Result<()> {
        let producer = self.producer.as_ref()
            .ok_or_else(|| anyhow!("生产者未初始化"))?;
        
        let message = ExecutionTaskMessage {
            header: crate::kafka::messages::MessageHeader::new(
                "execution_task",
                "koda-sandbox-manager"
            ),
            task_id,
            sandbox_id,
            execution_request: crate::kafka::messages::ExecutionRequest {
                source_code,
                language,
                entry_point,
                env_vars,
                stdin,
            },
            priority,
            timeout_ms,
            retry_count: 0,
            max_retries: 3,
        };
        
        producer.send_execution_task(&message).await
    }
    
    /// 发送系统指标
    pub async fn send_system_metrics(
        &self,
        metrics_type: crate::kafka::messages::MetricsType,
        metrics: std::collections::HashMap<String, f64>,
        labels: std::collections::HashMap<String, String>,
    ) -> Result<()> {
        let producer = self.producer.as_ref()
            .ok_or_else(|| anyhow!("生产者未初始化"))?;
        
        let message = SystemMetricsMessage {
            header: crate::kafka::messages::MessageHeader::new(
                "system_metrics",
                "koda-sandbox-manager"
            ),
            metrics_type,
            metrics,
            labels,
        };
        
        producer.send_system_metrics(&message).await
    }
    
    /// 发送错误报告
    pub async fn send_error_report(
        &self,
        level: ErrorLevel,
        error_code: String,
        error_message: String,
        error_details: Option<String>,
        sandbox_id: Option<SandboxId>,
        stack_trace: Option<String>,
        context: std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<()> {
        let producer = self.producer.as_ref()
            .ok_or_else(|| anyhow!("生产者未初始化"))?;
        
        let message = ErrorReportMessage {
            header: crate::kafka::messages::MessageHeader::new(
                "error_report",
                "koda-sandbox-manager"
            ),
            level,
            error_code,
            error_message,
            error_details,
            sandbox_id,
            stack_trace,
            context,
        };
        
        producer.send_error_report(&message).await
    }
    
    /// 发送健康检查结果
    pub async fn send_health_check(
        &self,
        component: String,
        status: HealthStatus,
        details: std::collections::HashMap<String, serde_json::Value>,
        response_time_ms: u64,
    ) -> Result<()> {
        let producer = self.producer.as_ref()
            .ok_or_else(|| anyhow!("生产者未初始化"))?;
        
        let message = HealthCheckMessage {
            header: crate::kafka::messages::MessageHeader::new(
                "health_check",
                "koda-sandbox-manager"
            ),
            component,
            status,
            details,
            response_time_ms,
        };
        
        producer.send_health_check(&message).await
    }
    
    /// 刷新生产者缓冲区
    pub async fn flush_producer(&self, timeout_ms: u64) -> Result<()> {
        if let Some(producer) = &self.producer {
            producer.flush(timeout_ms).await
        } else {
            Ok(())
        }
    }
    
    /// 获取配置
    pub fn config(&self) -> &KafkaConfig {
        &self.config
    }
}

/// 沙箱消息处理器
pub struct SandboxMessageHandler {
    sandbox_manager: Arc<SandboxManager>,
}

impl SandboxMessageHandler {
    pub fn new(sandbox_manager: Arc<SandboxManager>) -> Self {
        Self {
            sandbox_manager,
        }
    }
}

#[async_trait]
impl MessageHandler for SandboxMessageHandler {
    async fn handle_message(&self, message: KafkaMessage) -> Result<()> {
        debug!("处理Kafka消息: {}", message.message_type());
        
        match message {
            KafkaMessage::SandboxEvent(event) => {
                self.handle_sandbox_event(event).await
            }
            KafkaMessage::ExecutionTask(task) => {
                self.handle_execution_task(task).await
            }
            KafkaMessage::SystemMetrics(metrics) => {
                self.handle_system_metrics(metrics).await
            }
            KafkaMessage::ErrorReport(error) => {
                self.handle_error_report(error).await
            }
            KafkaMessage::HealthCheck(health) => {
                self.handle_health_check(health).await
            }
        }
    }
    
    async fn handle_error(&self, topic: &str, partition: i32, offset: i64, error: &str) {
        error!("Kafka消息处理错误 - Topic: {}, Partition: {}, Offset: {}, Error: {}", 
               topic, partition, offset, error);
    }
}

impl SandboxMessageHandler {
    /// 处理沙箱事件
    async fn handle_sandbox_event(&self, event: SandboxEventMessage) -> Result<()> {
        info!("处理沙箱事件: {:?} - 沙箱ID: {}", event.event_type, event.sandbox_id);
        
        // 可以在这里添加具体的事件处理逻辑
        // 例如：更新统计信息、触发告警等
        
        Ok(())
    }
    
    /// 处理执行任务
    async fn handle_execution_task(&self, task: ExecutionTaskMessage) -> Result<()> {
        info!("处理执行任务: {}", task.task_id);
        
        // 1. 分配或使用指定的沙箱
        let sandbox_id = if let Some(id) = task.sandbox_id {
            id
        } else {
            // 创建新沙箱
            let spec = crate::types::SandboxSpec {
                language: task.execution_request.language.clone(),
                cpu_cores: 1,
                memory_mb: 512,
                disk_mb: 1024,
                network_enabled: true,
                timeout_secs: task.timeout_ms / 1000,
                env_vars: task.execution_request.env_vars.clone(),
                snapshot_path: None,
                memory_file_path: None,
                kernel_image_path: None,
                rootfs_path: None,
            };
            
            self.sandbox_manager.create_sandbox(spec).await?
        };
        
        // 2. 执行代码
        let exec_request = crate::types::ExecutionRequest {
            sandbox_id,
            source_code: task.execution_request.source_code,
            language: task.execution_request.language,
            entry_point: task.execution_request.entry_point,
            dependencies: vec![], // 简化实现
            env_vars: task.execution_request.env_vars,
            stdin: task.execution_request.stdin,
            timeout_ms: Some(task.timeout_ms),
        };
        
        match self.sandbox_manager.execute_code(sandbox_id, exec_request).await {
            Ok(result) => {
                info!("执行任务完成: {} - 成功: {}", task.task_id, result.success);
                // 可以在这里发送执行结果到另一个主题
            }
            Err(e) => {
                error!("执行任务失败: {} - 错误: {}", task.task_id, e);
                // 可以在这里发送错误报告
            }
        }
        
        Ok(())
    }
    
    /// 处理系统指标
    async fn handle_system_metrics(&self, metrics: SystemMetricsMessage) -> Result<()> {
        debug!("处理系统指标: {:?} - {} 个指标", metrics.metrics_type, metrics.metrics.len());
        
        // 可以在这里将指标写入时序数据库或监控系统
        
        Ok(())
    }
    
    /// 处理错误报告
    async fn handle_error_report(&self, error: ErrorReportMessage) -> Result<()> {
        warn!("收到错误报告: {:?} - {}", error.level, error.error_message);
        
        // 可以在这里进行错误处理、告警等
        
        Ok(())
    }
    
    /// 处理健康检查
    async fn handle_health_check(&self, health: HealthCheckMessage) -> Result<()> {
        debug!("处理健康检查: {} - {:?}", health.component, health.status);
        
        // 可以在这里更新组件健康状态
        
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Config;

    #[tokio::test]
    async fn test_kafka_manager_creation() {
        let config = KafkaConfig::default();
        let manager = KafkaManager::new(config);
        
        assert!(manager.producer.is_none());
        assert!(manager.consumer.is_none());
    }

    #[test]
    fn test_sandbox_message_handler() {
        // 注意: 这个测试需要模拟的沙箱管理器
        // let sandbox_manager = Arc::new(SandboxManager::new(Arc::new(Config::default())).await.unwrap());
        // let handler = SandboxMessageHandler::new(sandbox_manager);
        // assert!(!handler.sandbox_manager.is_null());
    }
}